from base64 import b64decode
from email import message
from typing import override
import requests
from volcengine.visual.VisualService import VisualService
from urllib.parse import unquote
import logging

logger = logging.getLogger(__name__)


class JiMengClient():
    def __init__(self, ak: str=None, sk: str=None):
        self._client = VisualService()
        if ak:
           self._client.set_ak(ak)
        if sk:
            self._client.set_sk(sk)

    def generate_image(
            self,
            prompt: str,
            width: int = 512,
            height: int = 512,
            use_pre_llm: bool = False,
            use_sr: bool = True,
            return_url: bool = False,
            seed: int = -1) -> dict:
        body = self._create_body(
            prompt, width, height, use_pre_llm, use_sr, return_url, seed)

        response = self._client.cv_process(body)
        logger.info(f"generate image response: {response}")
        code = response["code"]
        is_success = code == 10000
        message = response["message"]
        request_id = response["request_id"]

        image = None
        image_type = None
        if is_success:
            image, image_type = self._retrieve_image_bytes(response)

        return Result({
            "is_success": is_success,
            "code": code,
            "message": message,
            "image": image,
            "request_id": request_id,
            "image_type": image_type
        })

    @property
    def access_key(self):
        return self._client.get_ak()

    @access_key.setter
    def access_key(self, value):
        self._client.set_ak(value)

    @property
    def secret_key(self):
        return self._client.get_sk()

    @secret_key.setter
    def secret_key(self, value):
        self._client.set_sk(value)

    def _create_body(self, prompt: str, width: int = 512, height: int = 512, use_pre_llm: bool = False, use_sr: bool = True, return_url: bool = True, seed: int = -1, req_key="jimeng_high_aes_general_v21_L") -> dict:
        return {
            "req_key": req_key,
            "prompt": prompt,
            "width": width,
            "height": height,
            "use_pre_llm": use_pre_llm,
            "use_sr": use_sr,
            "return_url": return_url,
            "seed": seed
        }

    def _retrieve_image_bytes(self, response):
        image = None
        image_type = None
        b64_image_list = response["data"]["binary_data_base64"]
        if b64_image_list and len(b64_image_list) > 0:
            image = b64decode(b64_image_list[0])
            image_type = 'base64'
        else:
            image_urls = response["data"]["image_urls"]
            if image_urls and len(image_urls) > 0:
                image_url = unquote(image_urls[0])
                image = self._download_image(image_url)
                image_type = 'image_url'

        return image, image_type

    def _download_image(self, image_url: str):
        logger.info(f"download image: {image_url}")
        response = requests.get(image_url)
        if response.status_code != 200:
            raise Exception(
                f"Failed to download image. url={image_url}, Status code: {response.status_code}, Message: {response.reason}.")
        logger.info(f"download image success: {image_url}")
        return response.content


class Result():
    def __init__(self, response):
        self.response = response

    @property
    def is_success(self):
        return self.response['is_success']

    @property
    def code(self):
        return self.response['code']

    @property
    def message(self):
        return self.response['message']

    @property
    def image(self):
        return self.response['image']

    @property
    def request_id(self):
        return self.response['request_id']

    @property
    def image_type(self):
        return self.response['image_type']

    @override
    def __str__(self) -> str:
        return f"Result(is_success={self.is_success}, code={self.code}, message={self.message}, request_id={self.request_id}, image_type={self.image_type}), image={'***' if self.image is not None else 'None'}"
